package com.zhx.common.common.aspect;


import com.zhx.common.common.annotation.RateLimit;
import com.zhx.common.common.exception.BusinessException;
import com.zhx.common.common.exception.CommonErrorEnum;
import com.zhx.common.common.service.rateLimit.*;
import com.zhx.common.common.util.AssertUtil;
import com.zhx.common.common.util.RequestHolder;
import jodd.util.StringUtil;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import utils.SpElUtils;

import java.lang.reflect.Method;

@Aspect
@Component
@Slf4j
public class RateLimitAspect {
    @Before("@annotation(com.zhx.common.common.annotation.RateLimit) || @annotation(com.zhx.common.common.annotation.RateLimits)")
    public void checkRateLimit(JoinPoint joinPoint) {
        log.info("checkRateLimit.....................");
        // 获取所有应用到方法上的RateLimit注解
        MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
        Method method = methodSignature.getMethod();

        RateLimit[] rateLimits = method.getAnnotationsByType(RateLimit.class);

        for (int i = 0; i < rateLimits.length; i++) {
            RateLimit limit = rateLimits[i];

            //生成key
            String prefix = StringUtil.isBlank(limit.prefixKey()) ? SpElUtils.getMethodKey(method) + ":index:" + i : limit.prefixKey();
            String key = "";
            // 根据注解指定的目标类型（EL、IP、UID），生成对应的键
            switch (limit.target()) {
                case EL:
                    key = SpElUtils.parseSpEl(method, joinPoint.getArgs(), limit.spEl());
                    break;
                case IP:
                    key = RequestHolder.get().getIp();
                    break;
                case UID:
                    key = RequestHolder.get().getUid().toString();
            }
            key = prefix + ":" + key;

            RateLimitStrategyHandler strategyHandler = getStrategy(limit.strategy());
            AssertUtil.isTrue(strategyHandler.shouldAllow(key,limit.count(), limit.time(),limit.unit()), CommonErrorEnum.LOCK_LIMIT);
        }
    }

    private RateLimitStrategyHandler getStrategy(RateLimitStrategy strategy) {
        switch (strategy) {
            case FIXED_WINDOW:
                return new FixedWindowStrategyHandler();
            case SLIDING_WINDOW:
                return new SlidingWindowStrategyHandler();
            case TOKEN_BUCKET:
                return new TokenBucketStrategyHandler();
            case LEAKY_BUCKET:
                return new LeakyBucketStrategyHandler();
            default:
                throw new BusinessException("Unsupported rate limit strategy: " + strategy);
        }
    }

}
